import json
import sys
import numpy as np
import matplotlib.pyplot as plt
import collections
import ast
import pickle
import multiprocessing
import os
from math import log
from math import exp
from heapq import nlargest
from time import sleep
from time import time, ctime
from itertools import groupby
from sklearn.pipeline import Pipeline
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import confusion_matrix
from sklearn.datasets import make_multilabel_classification
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.multiclass import OneVsRestClassifier
from sklearn.svm import LinearSVC
from sklearn.multiclass import OneVsOneClassifier
from sklearn import datasets
from sklearn.externals import joblib

from util import *

def kbeam(data, CH, PH, resultM, K=3):
	'''
	data is a list of Product() which is specified in util.py
	resultM is a ResultM() which is specified in util.py
	'''
	printLog(logFile, resultM.typeName)
	printLog(logFile, 'K='+str(K))
	analysisLog(resultM.typeName)
	analysisLog('K='+str(K))

	CLF_name = resultM.clf_name
	path = resultM.path

	initTime = time()

	parents = [0]

	resultM.pedningRoot(len(data))
	cnt = 0
	while len(parents)!= 0:
		children = []

		for p in parents:
			print 'processing', p, 'in', parents
			if (p in CH): 
				children += resultM.Her_tree[p].children
				cnt += 1
				# get the list of data and corresponding paths from the herairchical tree
				dataIdcs = []
				pathIdcs = []
				for productIdx, pathIdx in resultM.Her_tree[p].pendingProducts:
					dataIdcs.append(productIdx)
					pathIdcs.append(pathIdx)
				# reset the herarchical tree
				resultM.Her_tree[p].pendingProducts = []

				if len(dataIdcs) == 0:
					continue

				nodeTime = time()

				# get the resulting probability matrix and corresponding classes array
				classes = resultM.getResultM(p, dataIdcs, data)
				if type(classes) is int:
					continue

				# update data expended candidate pathes for each data in dataIdcs
				for dataIdx in range(len(dataIdcs)):
					for classIdx in range(len(classes)):
						tmpPath = Path()

						tmp = data[dataIdx]
						tmp = pathIdcs[dataIdx]
						tmp = classes[classIdx]
						tmp = classes[classIdx]
						tmp = data[dataIdx].predicted[pathIdcs[dataIdx]].path

						tmpPath.path = data[dataIdx].predicted[pathIdcs[dataIdx]].path+[classes[classIdx]]
						tmpPath.prob = data[dataIdx].predicted[pathIdcs[dataIdx]].prob+resultM.probMatrix[dataIdx][classIdx]

						data[dataIdx].candidates.append(tmpPath)

				# record processing performance
				nodeTime=time()-nodeTime
				analysisLog('#'+str(cnt)+' time ='+str(round((time()-initTime),1))+' ('+str(p)+')'+' size = '+str(len(dataIdcs))+' IT = '+str(nodeTime))

		# update predicted pathes for each data by selecting the K most probable candidates
		print parents, 'processed. Finalizing current round'
		parents = children
		for dataIdx in range(len(data)):
			proTime = time()
			candidates = data[dataIdx].candidates
			indices = range(len(candidates))
			if len(indices) != 0:
				tmpTime = time()
				largestKIdx = nlargest(K, indices, key=lambda j: candidates[j].prob)
				print 'largetest time',  time()-tmpTime
				data[dataIdx].predicted=[]
				for l in largestKIdx:
					pred = candidates[l].path[-1]
					resultM.Her_tree[pred].pendingProducts.append((dataIdx, len(data[dataIdx].predicted)))
					data[dataIdx].predicted.append(candidates[l])
				data[dataIdx].candidates=[]
				print 'assigning time', time()-tmpTime

			print 'data', dataIdx, 'finished in ', time()-proTime
		print 'new round starts'

	printLog(logFile,  'Testing is Done in '+ str(round(time()-time0)))

	# logging results
	predicted = []
	labels = []
	for dataIdx in range(len(data)):
		tmp = sorted(data[dataIdx].predicted, key = lambda x: x.prob, reverse=True)
		predicted.append(tmp[0].path)
		labels = data[dataIdx].actual

	resultLog(predicted, labels)

class ML_NB_ResultM(ResultM):
	'''
		clf_name stores the classifire name
		path stores the path to classifiers
		Her_tree stores the dict of nodes of the hierarchical tree
		probMatrix stores the resulting probability matrix after each getResultM() call
	'''

	def __init__(self, trainThres, PH, CH):
		self.clf_name = 'ML_NB'
		self.path = '../../../Data229/CLF/'+self.clf_name+str(trainThres)+'/'
		Her_tree = dict()
		# constructing the Her_tree
		parents = [0]
		built = []
		while parents!=[]:
			children = []
			for p in parents:
				if p in CH:
					children += CH[p]
				if p in built:
					continue
				built.append(p)
				if p in CH:
					cs = CH[p] 
				else:
					cs = []
				if p in PH:
					ps = PH[p]
				else:
					ps = []
				node = Node(0, p, cs, ps, 0)
				self.Her_tree[p]=node
			parents = children
		typeName = 'ML_NB_ResultM'


	def getResultM(self, p, dataIdcs, data):
		'''
			build the result matrix on node p, returns the classes array 
		'''
		dataContents = []
		for dataIdx in dataIdcs:
			dataContents.append(data[dataIdx].data)
		if self.Her_tree[p].baseClf == None:
			return 0

		if self.Her_tree[p].baseClf==0:
			if not os.path.isfile(self.path+str(p)+'.pkl'):
				clf = None
			else:
				clf = joblib.load(self.path+str(p)+'.pkl')
			self.Her_tree[p].baseClf=clf
		self.probMatrix = self.Her_tree[p].baseClf.predict_log_proba(dataContents)
		return self.Her_tree[p].baseClf.steps[-1][-1].classes_
		
